package com.example.demo.service.aiService;

import com.example.demo.demo.invoke.DialogueState;
import com.example.demo.model.dto.ai.AiDecision;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

@Service
public class StreamingAiService {

    private final ChatClient chatClient;
    private final ObjectMapper objectMapper;

    // --- 阶段一：决策调用的提示词 ---
    private static final String DECISION_PROMPT_TEMPLATE = """
            You are a dialogue policy engine. Your job is to analyze the dialogue state and user's query, update the state, and decide on the next action.
                    
            **Your internal thought process:**
            1.  **Update State:** Synthesize `Current Dialogue State` and `User's Latest Message` into an `updatedState`.
            2.  **Assess Sufficiency:** Look at the `updatedState`. Is there enough information to give a comprehensive final advice?
            3.  **Decide Action:** Based on your assessment, choose an action: `ASK_A_CLARIFYING_QUESTION` or `PROVIDE_FINAL_ADVICE`.

            **Your output MUST BE a single, valid JSON object** in the following format, with no other text.
            ```json
            \\{
              "action": "...",
              "updatedState": \\{ ... \\}
            \\}
            ```
                    
            ### Context ###
            **Current Dialogue State (JSON):**
            {currentStateJson}
                    
            **User's Latest Message:**
            "{latestQuery}"
            """;

    // --- 阶段二：生成调用的提示词 ---
    private static final String GENERATION_PROMPT_TEMPLATE = """
            You are "恋爱大师", a world-class empathetic AI relationship advisor.
            Your designated task for this turn is: **{actionToTake}**.
                    
            ### Context ###
            **Full Dialogue State:**
            {currentStateJson}
                    
            **User's Latest Message:**
            "{latestQuery}"
                    
            ### Your Task ###
            Based on your designated task and the full context, generate your response to the user now.
            - If your task is `ASK_A_CLARIFYING_QUESTION`, ask one warm, open-ended question.
            - If your task is `PROVIDE_FINAL_ADVICE`, provide the full, three-part structured response.
            """;

    public StreamingAiService(ChatClient.Builder chatClientBuilder, ObjectMapper objectMapper) {
        // 初始化基于内存的对话记忆
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(new InMemoryChatMemoryRepository())
                .maxMessages(10)
                .build();
        this.chatClient = chatClientBuilder
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build()
                )
                .build();
        this.objectMapper = objectMapper;
    }

    // 阶段一的实现方法 (非流式)
    public AiDecision getAiDecision(DialogueState currentState, String latestQuery) {
        String currentStateJson = convertStateToJson(currentState);
        return this.chatClient.prompt()
                .user(userSpec -> userSpec
                        .text(DECISION_PROMPT_TEMPLATE)
                        .param("currentStateJson", currentStateJson)
                        .param("latestQuery", latestQuery))
                .call()
                .entity(AiDecision.class);
    }

    // 阶段二的实现方法 (流式)
    public Flux<String> streamReply(AiDecision decision, String latestQuery) {
        String currentStateJson = convertStateToJson(decision.updatedState());
        return this.chatClient.prompt()
                .advisors(advisor -> advisor.param("chat_memory_conversation_id", "678")
                        .param("chat_memory_response_size", 100))
                .user(userSpec -> userSpec
                        .text(GENERATION_PROMPT_TEMPLATE)
                        .param("actionToTake", decision.action().toString())
                        .param("currentStateJson", currentStateJson)
                        .param("latestQuery", latestQuery))
                .stream() // <-- 关键点：调用 .stream() 而不是 .call()
                .content();
    }

    private String convertStateToJson(DialogueState state) {
        try {
            return objectMapper.writeValueAsString(state);
        } catch (JsonProcessingException e) {
            throw new RuntimeException("Error serializing DialogueState to JSON", e);
        }
    }
}